from typing import Any, Dict, List, Optional, Type
from pydantic import BaseModel, Field
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, ToolMessage


# 类型映射表：LangChain 类 <=> 字符串类型
_MESSAGE_TYPE_MAP: Dict[str, Type[BaseMessage]] = {
    "human": HumanMessage,
    "ai": AIMessage,
    "tool": ToolMessage,
}
_REVERSE_TYPE_MAP = {v: k for k, v in _MESSAGE_TYPE_MAP.items()}


class MessageModel(BaseModel):
    type: str = Field(..., description="Message type: human / ai / tool")
    content: str
    id: Optional[str] = None
    name: Optional[str] = None  # For ToolMessage
    tool_call_id: Optional[str] = None  # For ToolMessage
    tool_calls: Optional[Any] = None  # For AIMessage
    additional_kwargs: Dict[str, Any] = {}
    response_metadata: Dict[str, Any] = {}

    @classmethod
    def from_lc(cls, msg: BaseMessage) -> "MessageModel":
        return cls(
            type=_REVERSE_TYPE_MAP.get(type(msg), "unknown"),
            **msg.dict()
        )

    def to_lc(self) -> BaseMessage:
        msg_cls = _MESSAGE_TYPE_MAP.get(self.type)
        if not msg_cls:
            raise ValueError(f"Unknown message type: {self.type}")
        return msg_cls(**self.model_dump(exclude={"type"}))



# # ✅ 将消息对象序列化为 JSON
# models = [MessageModel.from_lc(m) for m in messages]
# json_str = json.dumps([m.model_dump() for m in models], indent=2, ensure_ascii=False)
# print(json_str)
